package main

import (
	"net"
	"os"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/mdlayher/raw"
	"github.com/pkg/errors"
	"golang.org/x/net/bpf"
)

// IPMACCacheTable 用于缓存IP-MAC地址的映射
var IPMACCacheTable = make(map[string]net.HardwareAddr)

// SendReceiveDHCP4Msg 发送和接收DHCP消息
func SendReceiveDHCP4Msg(inter *net.Interface, smac, dmac net.HardwareAddr, sip, dip net.IP, req *dhcpv4.DHCPv4, recvTimeout time.Duration) (*dhcpv4.DHCPv4, error) {
	// bpf rawInstruction 通过 tcpdump 命令获取：tcpdump -dd -i lo "udp dst port 68"
	sock, err := raw.ListenPacket(inter, 0x0800, &raw.Config{
		Filter: []bpf.RawInstruction{
			{0x28, 0, 0, 0x0000000c},
			{0x15, 0, 4, 0x000086dd},
			{0x30, 0, 0, 0x00000014},
			{0x15, 0, 11, 0x00000011},
			{0x28, 0, 0, 0x00000038},
			{0x15, 8, 9, 0x00000044},
			{0x15, 0, 8, 0x00000800},
			{0x30, 0, 0, 0x00000017},
			{0x15, 0, 6, 0x00000011},
			{0x28, 0, 0, 0x00000014},
			{0x45, 4, 0, 0x00001fff},
			{0xb1, 0, 0, 0x0000000e},
			{0x48, 0, 0, 0x00000010},
			{0x15, 0, 1, 0x00000044},
			{0x6, 0, 0, 0x00040000},
			{0x6, 0, 0, 0x00000000},
		},
	})
	if err != nil {
		return nil, errors.Wrap(err, "create a raw socket")
	}
	defer sock.Close()

	// 接受消息
	var (
		recvErrorCh = make(chan error)
		resp        *dhcpv4.DHCPv4
	)
	go func(errCh chan<- error) {
		buf := make([]byte, 4096)
		deadline := time.After(recvTimeout)
	FOR:
		for {
			select {
			case <-deadline:
				errCh <- errors.New("receiving DHCP response has timed out")
				return
			default:
				if err := sock.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
					errCh <- err
					return
				}

				n, _, err := sock.ReadFrom(buf)
				if err != nil {
					if os.IsTimeout(err) {
						continue FOR
					}
					errCh <- err
					return
				}

				var (
					smac net.HardwareAddr
					sip4 net.IP
				)

				// 判断报文
				pkt := gopacket.NewPacket(buf[:n], layers.LayerTypeEthernet, gopacket.Default)
				if ethLayer := pkt.Layer(layers.LayerTypeEthernet); ethLayer != nil {
					eth := ethLayer.(*layers.Ethernet)
					if eth.EthernetType != 0x0800 { // 不是IPv4
						continue FOR
					}
					smac = eth.SrcMAC
				}

				if ip4Layer := pkt.Layer(layers.LayerTypeIPv4); ip4Layer != nil {
					ip4 := ip4Layer.(*layers.IPv4)
					if ip4.Protocol != 17 { // 不是UDP协议
						continue FOR
					}
					sip4 = ip4.SrcIP
				}
				IPMACCacheTable[sip4.String()] = smac

				if udpLayer := pkt.Layer(layers.LayerTypeUDP); udpLayer != nil {
					udp := udpLayer.(*layers.UDP)
					if udp.SrcPort != dhcpv4.ServerPort && udp.DstPort != dhcpv4.ClientPort { // 端口不匹配
						continue FOR
					}
					var err error
					resp, err = dhcpv4.FromBytes(udp.Payload)
					if err != nil { // 不是正确的dhcp报文
						continue FOR
					}
				}

				// check that this is a response to our message
				if resp.TransactionID != req.TransactionID {
					continue FOR
				}
				// wait for a response message
				if resp.OpCode != dhcpv4.OpcodeBootReply {
					continue FOR
				}

				// 如果请求的类型是某个类型，应该收到的应答是指定的类型
				switch req.MessageType() {
				case dhcpv4.MessageTypeDiscover: // 应该收到offer类型的应答
					if resp.MessageType() != dhcpv4.MessageTypeOffer {
						continue FOR
					}
				case dhcpv4.MessageTypeRequest: // 应该收到ack，nack类型的应答
					if !(resp.MessageType() == dhcpv4.MessageTypeAck || resp.MessageType() == dhcpv4.MessageTypeNak) {
						continue FOR
					}
				default:
					errCh <- errors.New("message type of DHCP request is incorrect")
					return
				}

				errCh <- nil
				return
			}
		}
	}(recvErrorCh)

	// 发送消息
	ethFrameData, err := buildRawUDP4Packet(smac, dmac, sip, dip, dhcpv4.ClientPort, dhcpv4.ServerPort, req.ToBytes())
	if err != nil {
		return nil, errors.Wrap(err, "building a Ethernet data frame with DHCP message")
	}

	_, err = sock.WriteTo(ethFrameData, &raw.Addr{HardwareAddr: dmac})
	if err != nil {
		return nil, errors.Wrap(err, "failed to write to raw socket")
	}

	return resp, <-recvErrorCh
}

// SendDHCP4Msg 发送DHCP4消息
func SendDHCP4Msg(inter *net.Interface, smac, dmac net.HardwareAddr, sip, dip net.IP, req *dhcpv4.DHCPv4) error {
	sock, err := raw.ListenPacket(inter, 0x0800, &raw.Config{})
	if err != nil {
		return errors.Wrap(err, "create a raw socket")
	}
	defer sock.Close()

	ethFrameData, err := buildRawUDP4Packet(smac, dmac, sip, dip, dhcpv4.ClientPort, dhcpv4.ServerPort, req.ToBytes())
	if err != nil {
		return errors.Wrap(err, "building a Ethernet data frame with DHCP message")
	}
	_, err = sock.WriteTo(ethFrameData, &raw.Addr{HardwareAddr: dmac})
	return err
}

// getNetInterface 获取网络接口
func getNetInterface(ifName string) (*net.Interface, error) {
	ifs, err := net.Interfaces()
	if err != nil {
		return nil, errors.WithMessage(err, "get network interfaces list")
	}
	for _, itf := range ifs {
		if itf.Name == ifName {
			return &itf, nil
		}
	}

	return nil, errors.New("the specified network interface was not found")
}

// buildRawUDP4Packet 构建Ethernet-IP4-UDP数据包
func buildRawUDP4Packet(sMAC, dMAC net.HardwareAddr, sIP, dIP net.IP, sPort, dPort uint16, data []byte) ([]byte, error) {
	payload := gopacket.Payload(data)

	udpHeader := &layers.UDP{SrcPort: layers.UDPPort(sPort), DstPort: layers.UDPPort(dPort)}

	ip4Header := &layers.IPv4{
		Version:  4,
		Id:       uint16(os.Getpid() & 0xffff),
		TTL:      32,
		Protocol: 17,
		SrcIP:    sIP,
		DstIP:    dIP,
	}
	if err := udpHeader.SetNetworkLayerForChecksum(ip4Header); err != nil {
		return nil, errors.WithStack(err)
	}

	ethHeader := &layers.Ethernet{
		SrcMAC:       sMAC,
		DstMAC:       dMAC,
		EthernetType: 0x0800,
	}

	// 系列化
	var (
		buffer           = gopacket.NewSerializeBuffer()
		serializeOptions = gopacket.SerializeOptions{ComputeChecksums: true, FixLengths: true}
	)
	for _, layer := range []gopacket.SerializableLayer{payload, udpHeader, ip4Header, ethHeader} {
		err := layer.SerializeTo(buffer, serializeOptions)
		if err != nil {
			return nil, errors.Wrap(err, "failed to serialize dhcp package")
		}
		buffer.PushLayer(layer.LayerType())
	}

	return buffer.Bytes(), nil
}
